#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
模型下载脚本
用于下载Qwen1.5-1.8B模型到项目本地目录
"""

import os
from transformers import AutoModelForCausalLM, AutoTokenizer
import argparse

def download_model(model_name="Qwen/Qwen1.5-1.8B", cache_dir=None):
    """
    下载模型到本地项目目录
    
    Args:
        model_name (str): Hugging Face模型名称
        cache_dir (str): 指定缓存目录，如果为None则使用项目目录
    """
    # 如果没有指定缓存目录，则使用项目内的models目录
    if cache_dir is None:
        project_root = os.path.dirname(os.path.abspath(__file__))
        cache_dir = os.path.join(project_root, "models")
    
    model_dir = os.path.join(cache_dir, model_name.replace("/", "--"))
    print(f"开始下载模型: {model_name}")
    print(f"模型将被下载到: {model_dir}")
    
    # 确保缓存目录存在
    os.makedirs(cache_dir, exist_ok=True)
    
    try:
        # 下载模型和分词器
        print("正在下载模型文件...")
        model = AutoModelForCausalLM.from_pretrained(
            model_name,
            cache_dir=cache_dir,
            resume_download=True
        )
        print("正在下载分词器...")
        tokenizer = AutoTokenizer.from_pretrained(
            model_name,
            cache_dir=cache_dir,
            resume_download=True
        )
        
        print("模型下载完成!")
        print(f"模型位置: {model_dir}")
        
        # 显示模型信息
        if os.path.exists(model_dir):
            refs_path = os.path.join(model_dir, "refs", "main")
            if os.path.exists(refs_path):
                with open(refs_path, 'r') as f:
                    commit_hash = f.read().strip()
                    print(f"模型版本 (commit): {commit_hash}")
                    snapshot_path = os.path.join(model_dir, "snapshots", commit_hash)
                    if os.path.exists(snapshot_path):
                        print(f"模型快照路径: {snapshot_path}")
        
        return True
    except Exception as e:
        print(f"模型下载失败: {e}")
        return False

def main():
    parser = argparse.ArgumentParser(description="下载Qwen1.5-1.8B模型到项目本地目录")
    parser.add_argument(
        "--model-name", 
        default="Qwen/Qwen1.5-1.8B",
        help="要下载的Hugging Face模型名称 (默认: Qwen/Qwen1.5-1.8B)"
    )
    parser.add_argument(
        "--cache-dir",
        help="指定模型存储目录 (可选，默认为项目内models目录)"
    )
    
    args = parser.parse_args()
    
    success = download_model(args.model_name, args.cache_dir)
    if success:
        print("\n模型已成功下载到项目本地目录，现在可以离线使用。")
    else:
        print("\n模型下载失败，请检查网络连接或稍后重试。")

if __name__ == "__main__":
    main()